"""添加MCP相关表

Revision ID: 001_add_mcp_tables
Revises: 
Create Date: 2025-10-22

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import mysql

# revision identifiers, used by Alembic.
revision = '001_add_mcp_tables'
down_revision = None  # 根据实际情况修改为上一个迁移的revision ID
branch_labels = None
depends_on = None


def upgrade():
    """升级数据库：创建MCP相关表"""
    
    # 创建mcp_sessions表
    op.create_table(
        'mcp_sessions',
        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False, comment='主键'),
        sa.Column('user_id', sa.Integer(), nullable=False, comment='用户ID'),
        sa.Column('session_id', sa.String(128), nullable=False, comment='MCP会话ID（UUID）'),
        sa.Column('mcp_server_url', sa.String(512), nullable=False, comment='MCP服务器地址'),
        sa.Column('conversation_id', sa.Integer(), nullable=True, comment='关联的对话ID'),
        sa.Column('model_config_id', sa.Integer(), nullable=True, comment='关联的模型配置ID'),
        sa.Column('status', sa.Enum('connecting', 'connected', 'disconnected', 'error', name='sessionstatus'),
                  server_default='connecting', nullable=False, comment='会话状态'),
        sa.Column('is_active', sa.Boolean(), server_default='1', nullable=False, comment='是否活跃'),
        sa.Column('connection_config', sa.JSON(), nullable=True, comment='连接配置'),
        sa.Column('message_count', sa.Integer(), server_default='0', nullable=False, comment='消息总数'),
        sa.Column('error_count', sa.Integer(), server_default='0', nullable=False, comment='错误次数'),
        sa.Column('last_heartbeat_at', sa.DateTime(timezone=True), nullable=True, comment='最后心跳时间'),
        sa.Column('connected_at', sa.DateTime(timezone=True), nullable=True, comment='连接建立时间'),
        sa.Column('disconnected_at', sa.DateTime(timezone=True), nullable=True, comment='断开连接时间'),
        sa.Column('expires_at', sa.DateTime(timezone=True), nullable=True, comment='过期时间'),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('CURRENT_TIMESTAMP'),
                  nullable=False, comment='创建时间'),
        sa.Column('updated_at', sa.DateTime(timezone=True),
                  server_default=sa.text('CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP'),
                  nullable=False, comment='更新时间'),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('session_id'),
        sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE'),
        sa.ForeignKeyConstraint(['conversation_id'], ['conversations.id'], ondelete='SET NULL'),
        sa.ForeignKeyConstraint(['model_config_id'], ['model_configs.id'], ondelete='SET NULL'),
        mysql_charset='utf8mb4',
        mysql_engine='InnoDB',
        comment='MCP会话表'
    )
    
    # 创建索引
    op.create_index('idx_user_id', 'mcp_sessions', ['user_id'])
    op.create_index('idx_session_id', 'mcp_sessions', ['session_id'])
    op.create_index('idx_status', 'mcp_sessions', ['status'])
    op.create_index('idx_is_active', 'mcp_sessions', ['is_active'])
    op.create_index('idx_created_at_sessions', 'mcp_sessions', ['created_at'])
    
    # 创建mcp_contexts表
    op.create_table(
        'mcp_contexts',
        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False, comment='主键'),
        sa.Column('session_id', sa.String(128), nullable=False, comment='MCP会话ID'),
        sa.Column('model_id', sa.String(128), nullable=True, comment='模型ID'),
        sa.Column('context_data', sa.JSON(), nullable=False, comment='上下文数据'),
        sa.Column('context_type', sa.Enum('system', 'user', 'assistant', 'tool', name='contexttype'),
                  server_default='user', nullable=False, comment='上下文类型'),
        sa.Column('context_metadata', sa.JSON(), nullable=True, comment='元数据'),
        sa.Column('version', sa.Integer(), server_default='1', nullable=False, comment='版本号'),
        sa.Column('token_count', sa.Integer(), server_default='0', nullable=False, comment='Token数量'),
        sa.Column('size_bytes', sa.Integer(), server_default='0', nullable=False, comment='数据大小（字节）'),
        sa.Column('is_encrypted', sa.Boolean(), server_default='0', nullable=False, comment='是否加密'),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('CURRENT_TIMESTAMP'),
                  nullable=False, comment='创建时间'),
        sa.Column('updated_at', sa.DateTime(timezone=True),
                  server_default=sa.text('CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP'),
                  nullable=False, comment='更新时间'),
        sa.PrimaryKeyConstraint('id'),
        sa.ForeignKeyConstraint(['session_id'], ['mcp_sessions.session_id'], ondelete='CASCADE'),
        mysql_charset='utf8mb4',
        mysql_engine='InnoDB',
        comment='MCP上下文表'
    )
    
    # 创建索引
    op.create_index('idx_session_id_ctx', 'mcp_contexts', ['session_id'])
    op.create_index('idx_context_type', 'mcp_contexts', ['context_type'])
    op.create_index('idx_created_at_ctx', 'mcp_contexts', ['created_at'])
    
    # 创建mcp_messages表
    op.create_table(
        'mcp_messages',
        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False, comment='主键'),
        sa.Column('session_id', sa.String(128), nullable=False, comment='MCP会话ID'),
        sa.Column('message_id', sa.String(128), nullable=False, comment='消息ID（UUID）'),
        sa.Column('message_type', sa.Enum('CONNECT', 'CONNECT_ACK', 'REQUEST', 'RESPONSE',
                                          'ERROR', 'PING', 'PONG', 'DISCONNECT', name='messagetype'),
                  nullable=False, comment='消息类型'),
        sa.Column('direction', sa.Enum('send', 'receive', name='messagedirection'),
                  nullable=False, comment='消息方向'),
        sa.Column('payload', sa.JSON(), nullable=False, comment='消息负载'),
        sa.Column('status', sa.Enum('pending', 'sent', 'received', 'failed', name='messagestatus'),
                  server_default='pending', nullable=False, comment='消息状态'),
        sa.Column('error_message', sa.Text(), nullable=True, comment='错误信息'),
        sa.Column('send_at', sa.DateTime(timezone=True), nullable=True, comment='发送时间'),
        sa.Column('receive_at', sa.DateTime(timezone=True), nullable=True, comment='接收时间'),
        sa.Column('duration_ms', sa.Integer(), nullable=True, comment='处理时长（毫秒）'),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('CURRENT_TIMESTAMP'),
                  nullable=False, comment='创建时间'),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('message_id'),
        sa.ForeignKeyConstraint(['session_id'], ['mcp_sessions.session_id'], ondelete='CASCADE'),
        mysql_charset='utf8mb4',
        mysql_engine='InnoDB',
        comment='MCP消息表'
    )
    
    # 创建索引
    op.create_index('idx_session_id_msg', 'mcp_messages', ['session_id'])
    op.create_index('idx_message_id', 'mcp_messages', ['message_id'])
    op.create_index('idx_message_type', 'mcp_messages', ['message_type'])
    op.create_index('idx_status_msg', 'mcp_messages', ['status'])
    op.create_index('idx_created_at_msg', 'mcp_messages', ['created_at'])


def downgrade():
    """降级数据库：删除MCP相关表"""
    
    # 删除表（按依赖关系逆序删除）
    op.drop_table('mcp_messages')
    op.drop_table('mcp_contexts')
    op.drop_table('mcp_sessions')
    
    # 删除Enum类型
    op.execute('DROP TYPE IF EXISTS sessionstatus')
    op.execute('DROP TYPE IF EXISTS contexttype')
    op.execute('DROP TYPE IF EXISTS messagetype')
    op.execute('DROP TYPE IF EXISTS messagedirection')
    op.execute('DROP TYPE IF EXISTS messagestatus')

